package com.yifeng.repo.micro.service.server.engine.dal.dsx;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.processor.DsProcessor;
import com.baomidou.dynamic.datasource.support.DataSourceClassResolver;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.yifeng.repo.micro.service.server.engine.dal.TransactionTargetManager;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;

/**
 * 分省动态数据源注解拦截器
 *
 * @see com.baomidou.dynamic.datasource.aop.DynamicDataSourceAnnotationInterceptor
 * Created by daibing on 2024/6/18.
 */
public class ProvinceDynamicDataSourceAnnotationInterceptor implements MethodInterceptor {

    /**
     * The identification of SPEL.
     */
    private static final String DYNAMIC_PREFIX = "#";
    private static final String MASTER = "master";

    private final DataSourceClassResolver dataSourceClassResolver;
    private final DsProcessor dsProcessor;

    public ProvinceDynamicDataSourceAnnotationInterceptor(Boolean allowedPublicOnly, DsProcessor dsProcessor) {
        dataSourceClassResolver = new ProvinceDataSourceClassResolver(allowedPublicOnly);
        this.dsProcessor = dsProcessor;
    }

    /**
     * @see DynamicRoutingDataSource#determineDataSource
     * 将根据这里选择的dsKey去获取数据源
     */
    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        // master方法直接走主库，不再切换数据源
        String methodName = invocation.getMethod().getName();
        if (methodName.startsWith(MASTER)) {
            return invocation.proceed();
        }

        // 写事务走主库，不再切换数据源
        Class<?> clazz = invocation.getMethod().getDeclaringClass();
        String target = clazz.getName() + "." + methodName;
        boolean writeMethod = TransactionTargetManager.containsWriteMethod(target, true);
        if (writeMethod) {
            return invocation.proceed();
        }

        // 非master的读事务，支持切换数据源，继续动态数据源的AOP切换逻辑
        String dsKey = determineDatasourceKey(invocation);
        DynamicDataSourceContextHolder.push(dsKey);
        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }

    private String determineDatasourceKey(MethodInvocation invocation) {
        String key = dataSourceClassResolver.findKey(invocation.getMethod(), invocation.getThis());
        return key.startsWith(DYNAMIC_PREFIX) ? dsProcessor.determineDatasource(invocation, key) : key;
    }
}
